#!/usr/bin/env python
# -*- coding: utf-8 -*-
#  thirdorder, help compute anharmonic IFCs from minimal sets of displacements
#  Copyright (C) 2012-2018 Wu Li <wu.li.phys2011@gmail.com>
#  Copyright (C) 2012-2018 Jesús Carrete Montaña <jcarrete@gmail.com>
#  Copyright (C) 2012-2018 Natalio Mingo Bisquert <natalio.mingo@cea.fr>
#  Copyright (C) 2014-2018 Antti J. Karttunen <antti.j.karttunen@iki.fi>
#  Copyright (C) 2016-2018 Genadi Naydenov <gan503@york.ac.uk>
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
from os import listdir,system
import sys
import os.path

try:
    import cStringIO as StringIO
except ImportError:
    try:
        import StringIO
    except ImportError:
        import io as StringIO
try:
    import hashlib
    hashes = True
except ImportError:
    hashes = False

import thirdorder_core
from thirdorder_common import *
from ase.io import read
from ase import Atoms
from ase.io.trajectory import TrajectoryWriter
from ase.calculators.singlepoint import SinglePointCalculator
# from irff.md.gulp import get_gulp_forces

def write_gulp_in(A, runword='gradient qiterative nosymmetry conv verb debu',
                  time_step=0.1,
                  tot_step=10.0,
                  maxcyc=2000,
                  output=None,
                  restart=100,
                  lib='reax',movieFrequency=10,**kwargs):
    ''' runword = keyword in gulp input
        can be 'md conv' 'md conp' 'opti conp'
        qite = iterative solve charges, norx = no charges '''
    finp = open('inp-gulp', 'w')
    rw = runword.split()
    print(runword, file=finp)
    print('#', file=finp)
    print('#', file=finp)
    if rw[0] == 'opti':
        print('maxcyc %d' % maxcyc, file=finp)
        if 'gtol' in kwargs:
            print('gtol opt  {:f} '.format(kwargs['gtol']), file=finp)
        else:
            print('gtol opt  0.00001', file=finp)
    elif rw[0] == 'md':
        #  if rw[1]=='conv':
        #     print('ensemble nvt',file=finp)
        #     print('tau_thermostat     0.05 ps',file=finp)
        #  if rw[1]=='conp':
        #     print('integrator leapfrog verlet',file=finp)
        #     print('ensemble           npt',file=finp)
        #     print('pressure           0.00 GPa',file=finp)
        #     print('tau_barostat       0.1 ps',file=finp)
        #     print('tau_thermostat     0.1 ps',file=finp)
        if 'ensemble' in kwargs:
            print('ensemble   =  {:s} '.format(kwargs['ensemble']), file=finp)
        if 'tau_thermostat' in kwargs:
            print('tau_thermostat   =  {:f} ps'.format(kwargs['tau_thermostat']), file=finp)

        # print('temperature        %f K' % T, file=finp)
        print('timestep           %f fs' % time_step, file=finp)
        print('production         %f ps' %
              float(tot_step*time_step/1000.0), file=finp)
        print('equilibration      0.0 ps', file=finp)
        print('write              1', file=finp)
        print('sample             1', file=finp)
        
    if 'T' in kwargs: 
       print('temperature          {:f} K'.format(kwargs['T']),file=finp)
    if 'pressure' in kwargs:
       print('pressure          {:f} GPa'.format(kwargs['pressure']),file=finp)

    print('#', file=finp)
    print('title', file=finp)
    print('GULP calculation', file=finp)
    print('end', file=finp)
    print('#', file=finp)
    if lib == 'brenner':
        print('{:s}'.format(lib), file=finp)
    else:
        print('library %s' % lib, file=finp)

    if output is None:
        print('output movie {:d} xyz his.xyz'.format(
            movieFrequency), file=finp)
    else:
        print('output {:s}'.format(output), file=finp)
    print('output cif gulp produces gulp.cif', file=finp)
    print('dump {:d} restart.grs'.format(restart), file=finp)
    print('#', file=finp)

    print('vectors', file=finp)
    cell = A.get_cell()
    for a in cell:
        print(a[0], a[1], a[2], file=finp)
    print('\n', file=finp)

    print('cartesian', file=finp)
    pos = A.get_positions()
    symb = A.get_chemical_symbols()
    for sp, x in zip(symb, pos):
        print(sp, 'core', x[0], x[1], x[2], 0.0, 1.0, 0.0, 1, 1, 1, file=finp)
    finp.close()
    
def get_gulp_force_from_drv(fdrv, wforce=False):
    forces = []
    with open(fdrv, 'r') as fdrv:
        gl = -1
        for i, line in enumerate(fdrv.readlines()):
            if line.find('gradients cartesian eV/Ang') >= 0:
                gl = i
                natom = int(line.split()[3])
            if gl > 0:
                if i > gl and i <= gl+natom:
                    forces.append([-float(line.split()[1]), -float(line.split()[2]),
                                   -float(line.split()[3])])
    # print(forces)
    if wforce:
        with open('Forces.FA', 'w') as fsiesta:
            print(natom, file=fsiesta)
            for i, f in enumerate(forces):
                print('{:4d} {:12.8f} {:12.8f} {:12.8f}'.format(
                    i+1, f[0], f[1], f[2]), file=fsiesta)  # Why?
                # print('{:4d} {:12.8f} {:12.8f} {:12.8f}'.format(i+1,f[0],f[1],f[2]),file=fsiesta)
    return np.array(forces)
    
def get_reax_energy(fo='out',e_kw='ReaxFF force field'):
    fout = open(fo, 'r')
    e=0.0
    ebond=0.0
    elp=0.0
    eover=0.0
    eunder=0.0 
    eang=0.0 
    epen=0.0 
    tconj=0.0 
    etor=0.0 
    fconj=0.0 
    evdw=0.0 
    ehb=0.0 
    ecl=0.0 
    esl = 0.0
    for line in fout.readlines():
        if line.find('E(bond)') >= 0:
            ebond = float(line.split()[2])
        elif line.find('E(lonepair)') >= 0:
            elp = float(line.split()[2])
        elif line.find('E(over)') >= 0:
            eover = float(line.split()[2])
        elif line.find('E(under)') >= 0:
            eunder = float(line.split()[2])
        elif line.find('E(val)') >= 0:
            eang = float(line.split()[2])
        elif line.find('E(coa)') >= 0:
            tconj = float(line.split()[2])
        elif line.find('E(pen)') >= 0:
            epen = float(line.split()[2])
        elif line.find('E(tors)') >= 0:
            etor = float(line.split()[2])
        elif line.find('E(conj)') >= 0:
            fconj = float(line.split()[2])
        elif line.find('E(vdw)') >= 0:
            evdw = float(line.split()[2])
        elif line.find('E(hb)') >= 0:
            ehb = float(line.split()[2])
        elif line.find('E(coulomb)') >= 0:
            ecl = float(line.split()[2])
        elif line.find('E(self)') >= 0:
            esl = float(line.split()[2])
        elif line.find(e_kw) >= 0:
            e = float(line.split()[-2])
    fout.close()
    return e, ebond, elp, eover, eunder, eang, epen, tconj, etor, fconj, evdw, ehb, ecl, esl

def get_gulp_forces(images, traj='gulp_force.traj', ffield='reaxff_nn',e_kw='ReaxFF force field', wforce=False):
    ''' calculate the force by GULP '''
    his = TrajectoryWriter(traj, mode='w')
    for atoms in images:
        write_gulp_in(atoms, runword='gradient nosymmetry conv qite verb',
                      output='drv gulp', restart=1,
                      lib=ffield)
        system('gulp<inp-gulp>out')
        forces = get_gulp_force_from_drv('gulp.drv', wforce=wforce)
        #atoms.forces = forces
        (e, eb_, el_, eo_, eu_, ea_, ep_,
         etc_, et_, ef_, ev_, ehb_, ecl_, esl_) = get_reax_energy(fo='out',e_kw=e_kw)
        #e = atoms.get_potential_energy()
        atoms.calc = SinglePointCalculator(atoms, energy=e,forces=forces)
        # atoms.set_calculator(calc)
        his.write(atoms=atoms)
    his.close()
    return atoms

def read_POSCAR(directory):
    """
    Return all the relevant information contained in a POSCAR file.
    """
    with dir_context(directory):
        nruter = dict()
        nruter["lattvec"] = np.empty((3, 3))
        f = open("POSCAR.0", "r")
        firstline = next(f)
        factor = .1 * float(next(f).strip())
        for i in range(3):
            nruter["lattvec"][:, i] = [float(j) for j in next(f).split()]
        nruter["lattvec"] *= factor
        line = next(f)
        fields = next(f).split()
        old = False
        try:
            int(fields[0])
        except ValueError:
            old = True
        if old:
            nruter["elements"] = firstline.split()
            nruter["numbers"] = np.array([int(i) for i in line.split()])
            typeline = "".join(fields)
        else:
            nruter["elements"] = line.split()
            nruter["numbers"] = np.array(
                [int(i) for i in fields], dtype=np.intc)
            typeline = next(f)
        natoms = nruter["numbers"].sum()
        nruter["positions"] = np.empty((3, natoms))
        for i in range(natoms):
            nruter["positions"][:, i] = [float(j) for j in next(f).split()]
        f.close()
    nruter["types"] = []
    for i in range(len(nruter["numbers"])):
        nruter["types"] += [i] * nruter["numbers"][i]
    if typeline[0] == "C":
        nruter["positions"] = sp.linalg.solve(nruter["lattvec"],
                                              nruter["positions"] * factor)
    return nruter


def write_POSCAR(poscar, filename):
    """
    Write the contents of poscar to filename.
    """
    global hashes
    f = StringIO.StringIO()
    f.write("1.0\n")
    for i in range(3):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format((
            poscar["lattvec"][:, i] * 10.).tolist()))
    f.write("{0}\n".format(" ".join(poscar["elements"])))
    f.write("{0}\n".format(" ".join([str(i) for i in poscar["numbers"]])))
    f.write("Direct\n")
    for i in range(poscar["positions"].shape[1]):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format(
            poscar["positions"][:, i].tolist()))
    if hashes:
        toencode = f.getvalue()
        if sys.hexversion >= 0x3000000:
            toencode = toencode.encode("utf-8")
        header = hashlib.sha1(toencode).hexdigest()
    else:
        header = filename
    with open(filename, "w") as finalf:
        finalf.write("{0}\n".format(header))
        finalf.write(f.getvalue())
    f.close()


def normalize_SPOSCAR(sposcar):
    """
    Rearrange sposcar, as generated by gen_SPOSCAR, so that it is in
    valid VASP order, and return the result.
    """
    nruter = copy.deepcopy(sposcar)
    # Order used internally (from most to least significant):
    # k,j,i,iat For VASP, iat must be the most significant index,
    # i.e., atoms of the same element must go together.
    indices = np.array(range(nruter["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten().tolist()
    nruter["positions"] = nruter["positions"][:, indices]
    nruter["types"].sort()
    return nruter


def read_forces(filename):
    """
    Read a set of forces on atoms from filename, presumably in
    vasprun.xml format.
    """
    calculation = ElementTree.parse(filename).getroot().find("calculation")
    for a in calculation.findall("varray"):
        if a.attrib["name"] == "forces":
            break
    nruter = []
    for i in a.getchildren():
        nruter.append([float(j) for j in i.text.split()])
    nruter = np.array(nruter, dtype=np.double)
    return nruter


def build_unpermutation(sposcar):
    """
    Return a list of integers mapping the atoms in the normalized
    version of sposcar to their original indices.
    """
    indices = np.array(range(sposcar["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten()
    return indices.argsort().tolist()


if __name__ == "__main__":
    ''' thirdorder_gulp.py sow|reap na nb nc cutoff [nm/-integer] 
        thirdorder_gulp.py sow 4 4 1 -3  # create input file
        thirdorder_gulp.py sow 4 4 1 10  # create input file
        thirdorder_gulp.py sow 4 4 1 -3  # create force constant 3rd 
    '''
    if len(sys.argv) != 6 or sys.argv[1] not in ("sow", "reap"):
        sys.exit("Usage: {0} sow|reap na nb nc cutoff[nm/-integer]".format(
            sys.argv[0]))
    action = sys.argv[1]
    na, nb, nc = [int(i) for i in sys.argv[2:5]]

    if min(na, nb, nc) < 1:
        sys.exit("Error: na, nb and nc must be positive integers")
    if sys.argv[5][0] == "-":
        try:
            nneigh = -int(sys.argv[5])
        except ValueError:
            sys.exit("Error: invalid cutoff")
        if nneigh == 0:
            sys.exit("Error: invalid cutoff")
    else:
        nneigh = None
        try:
            frange = float(sys.argv[5])
        except ValueError:
            sys.exit("Error: invalid cutoff")
        if frange == 0.:
            sys.exit("Error: invalid cutoff")      
#   - end - parse argument options - end - parse argument options - end - parse argument options

    print("Reading POSCAR.0")
    poscar = read_POSCAR(".")
    natoms = len(poscar["types"])
    print("Analyzing the symmetries")
    symops = thirdorder_core.SymmetryOperations(
        poscar["lattvec"], poscar["types"], poscar["positions"].T, SYMPREC)
    print("- Symmetry group {0} detected".format(symops.symbol))
    print("- {0} symmetry operations".format(symops.translations.shape[0]))
    print("Creating the supercell")
    sposcar = gen_SPOSCAR(poscar, na, nb, nc)
    ntot = natoms * na * nb * nc
    print("Computing all distances in the supercell")
    dmin, nequi, shifts = calc_dists(sposcar)
    if nneigh != None:
       frange = calc_frange(poscar, sposcar, nneigh, dmin)
       print("- Automatic cutoff: {0} nm".format(frange))
    else:
       print("- User-defined cutoff: {0} nm".format(frange))

    print("Looking for an irreducible set of third-order IFCs")
    wedge = thirdorder_core.Wedge(poscar, sposcar, symops, dmin, nequi, shifts,
                                  frange)
    print("- {0} triplet equivalence classes found".format(wedge.nlist))
    list4 = wedge.build_list4()
    nirred = len(list4)
    nruns = 4 * nirred
    print("- {0} DFT runs are needed".format(nruns))
    if action == "sow":
        print(sowblock)
        print("Writing undisplaced coordinates to 3RD.SPOSCAR")
        write_POSCAR(normalize_SPOSCAR(sposcar), "3RD.SPOSCAR")
        width = len(str(4 * (len(list4) + 1)))
        namepattern = "3RD.POSCAR.{{0:0{0}d}}".format(width)
        print("Writing displaced coordinates to 3RD.POSCAR.*")
        for i, e in enumerate(list4):
            for n in range(4):
                isign = (-1)**(n // 2)
                jsign = -(-1)**(n % 2)
                # Start numbering the files at 1 for aesthetic
                # reasons.
                number = nirred * n + i + 1
                dsposcar = normalize_SPOSCAR(
                    move_two_atoms(sposcar, e[1], e[3], isign * H, e[0], e[2],
                                   jsign * H))
                filename = namepattern.format(number)
                write_POSCAR(dsposcar, filename)
    else: 
#     --- generating forces - generating forces - generating forces - generating forces ---
        # try: 
        #    system('mv POSCAR POSCAR.0')
        # except:
        #    print(' ')
        print(reapblock)                        # compute force constant
        # print("XML ElementTree implementation: {0}".format(xmllib))
        # print("Waiting for a list of vasprun.xml files on stdin")
        cars = listdir()
        filelist = []
        ind_     = []
        for car in cars:
           if car.find('3RD.POSCAR')>=0:
              # system('mv '+ car + 'POSCAR')
              ind_.append(int(car.split('.')[2]))
              filelist.append(car)

        ind = np.argsort(ind_)
        filelist = np.array(filelist)[ind]
        # atoms = read('POSCAR')
        # atoms = get_gulp_forces([atoms])
        nfiles = len(filelist)
        print("- {0} filenames read".format(nfiles))

        if nfiles != nruns:
           sys.exit("Error: {0} filenames were expected".format(nruns))
        for i in filelist:
            if not os.path.isfile(i):
               sys.exit("Error: {0} is not a regular file".format(i))
        print("Reading the forces")

        p = build_unpermutation(sposcar)
        
        forces = []
        for i in filelist:
            # forces.append(read_forces(i)[p, :])
            system('cp '+ i + ' POSCAR')
            atoms = read('POSCAR')
            #atoms = get_gulp_forces([atoms])
            atoms = get_gulp_forces([atoms],ffield='brenner',e_kw='Brenner potentials')
            forces.append(atoms.get_forces()[p, :])
            print("- {0} read successfully".format(i))
            res = forces[-1].mean(axis=0)
            print("- \t Average force:")
            print("- \t {0} eV/(A * atom)".format(res))
        print("Computing an irreducible set of anharmonic force constants")
        
        phipart = np.zeros((3, nirred, ntot))
        for i, e in enumerate(list4):
            for n in range(4):
                isign = (-1)**(n // 2)
                jsign = -(-1)**(n % 2)
                number = nirred * n + i
                phipart[:, i, :] -= isign * jsign * forces[number].T
        phipart /= (400. * H * H)
        print("Reconstructing the full array")
        phifull = thirdorder_core.reconstruct_ifcs(phipart, wedge, list4,
                                                   poscar, sposcar)
        print("Writing the constants to FORCE_CONSTANTS_3RD")
        write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange,
                   "FORCE_CONSTANTS_3RD")
    print(doneblock)
